#!/usr/bin/env python3
"""
RocketMQ 生产级消息发送者

功能特性：
- 连接池管理和自动恢复
- 智能重试机制
- 性能监控和指标收集
- 异步和同步发送支持
- 消息压缩和批量发送
- 熔断器模式
- 健康检查

Author: Python RocketMQ Guide
Date: 2025-08-18
"""

import asyncio
import json
import logging
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum
from typing import Dict, Any, Optional, List, Union, Callable
import queue
import weakref

try:
    from rocketmq.client import Producer, Message, SendResult
    from rocketmq.client import ConsumeStatus, PullConsumer
except ImportError:
    raise ImportError("请安装 rocketmq-client-python: pip install rocketmq-client-python")

from threading import Lock, Event
import random
import math

# 假设我们已有配置管理模块
# from config_manager import get_config, RocketMQConfig

logger = logging.getLogger(__name__)


class MessagePriority(Enum):
    """消息优先级"""
    LOW = 0
    NORMAL = 1
    HIGH = 2
    CRITICAL = 3


@dataclass
class MessageContext:
    """消息上下文信息"""
    message_id: str
    topic: str
    tags: Optional[str]
    keys: Optional[str]
    body: Union[str, bytes, dict]
    properties: Optional[Dict[str, str]]
    priority: MessagePriority = MessagePriority.NORMAL
    delay_level: int = 0
    created_at: float = None
    retry_count: int = 0
    max_retries: int = 3

    def __post_init__(self):
        if self.created_at is None:
            self.created_at = time.time()
        if isinstance(self.body, dict):
            self.body = json.dumps(self.body, ensure_ascii=False)


class SendStatus(Enum):
    """发送状态"""
    SUCCESS = "SUCCESS"
    FAILED = "FAILED"
    TIMEOUT = "TIMEOUT"
    CIRCUIT_BREAKER_OPEN = "CIRCUIT_BREAKER_OPEN"


@dataclass
class SendResponse:
    """发送响应"""
    status: SendStatus
    message_id: str
    queue_id: Optional[int] = None
    queue_offset: Optional[int] = None
    error: Optional[Exception] = None
    retry_count: int = 0
    elapsed_time: float = 0
    broker_name: Optional[str] = None


class CircuitBreaker:
    """熔断器实现"""

    def __init__(self, failure_threshold: int = 5, timeout: int = 60,
                 recovery_timeout: int = 30):
        """
        初始化熔断器

        Args:
            failure_threshold: 失败次数阈值
            timeout: 熔断超时时间（秒）
            recovery_timeout: 恢复检测超时时间（秒）
        """
        self.failure_threshold = failure_threshold
        self.timeout = timeout
        self.recovery_timeout = recovery_timeout
        self.failure_count = 0
        self.last_failure_time = None
        self.state = "CLOSED"  # CLOSED, OPEN, HALF_OPEN
        self._lock = Lock()

    def can_execute(self) -> bool:
        """检查是否可以执行操作"""
        with self._lock:
            if self.state == "CLOSED":
                return True
            elif self.state == "OPEN":
                if time.time() - self.last_failure_time > self.recovery_timeout:
                    self.state = "HALF_OPEN"
                    logger.info("熔断器切换到半开状态")
                    return True
                return False
            else:  # HALF_OPEN
                return True

    def record_success(self):
        """记录成功"""
        with self._lock:
            self.failure_count = 0
            if self.state == "HALF_OPEN":
                self.state = "CLOSED"
                logger.info("熔断器恢复到关闭状态")

    def record_failure(self):
        """记录失败"""
        with self._lock:
            self.failure_count += 1
            self.last_failure_time = time.time()

            if self.failure_count >= self.failure_threshold:
                if self.state != "OPEN":
                    self.state = "OPEN"
                    logger.warning(f"熔断器开启，失败次数: {self.failure_count}")


class ConnectionPool:
    """RocketMQ 生产者连接池"""

    def __init__(self, config, pool_size: int = 5):
        """
        初始化连接池

        Args:
            config: RocketMQ 配置
            pool_size: 连接池大小
        """
        self.config = config
        self.pool_size = pool_size
        self._pool = queue.Queue(maxsize=pool_size)
        self._active_connections = weakref.WeakSet()
        self._lock = Lock()
        self._created_count = 0
        self._closed = False

        # 预创建连接
        self._initialize_pool()

    def _initialize_pool(self):
        """初始化连接池"""
        for _ in range(self.pool_size):
            try:
                producer = self._create_producer()
                self._pool.put(producer)
            except Exception as e:
                logger.error(f"创建生产者连接失败: {e}")

    def _create_producer(self) -> Producer:
        """创建生产者实例"""
        try:
            producer = Producer(self.config.producer_group)
            producer.set_name_server_address(self.config.name_server_address)

            # 设置生产者参数
            producer.set_max_message_size(self.config.producer_max_message_size)
            producer.set_send_msg_timeout(self.config.producer_send_timeout)
            producer.set_retry_times_when_send_failed(self.config.producer_retry_times)

            # 如果有认证信息
            if self.config.access_key and self.config.secret_key:
                producer.set_session_credentials(
                    self.config.access_key,
                    self.config.secret_key,
                    ""
                )

            producer.start()

            with self._lock:
                self._created_count += 1

            logger.debug(f"创建新的生产者连接，总数: {self._created_count}")
            return producer

        except Exception as e:
            logger.error(f"创建生产者失败: {e}")
            raise

    @contextmanager
    def get_connection(self):
        """获取连接（上下文管理器）"""
        if self._closed:
            raise RuntimeError("连接池已关闭")

        producer = None
        try:
            # 尝试从池中获取连接
            try:
                producer = self._pool.get(timeout=5)
            except queue.Empty:
                # 池中没有可用连接，创建新的
                logger.info("连接池为空，创建新连接")
                producer = self._create_producer()

            # 检查连接是否有效
            if not self._is_connection_valid(producer):
                logger.warning("连接无效，重新创建")
                try:
                    producer.shutdown()
                except:
                    pass
                producer = self._create_producer()

            self._active_connections.add(producer)
            yield producer

        except Exception as e:
            logger.error(f"获取连接失败: {e}")
            raise
        finally:
            # 归还连接到池中
            if producer and not self._closed:
                try:
                    self._active_connections.discard(producer)
                    if self._is_connection_valid(producer):
                        self._pool.put(producer, timeout=1)
                    else:
                        producer.shutdown()
                except Exception as e:
                    logger.error(f"归还连接失败: {e}")
                    try:
                        producer.shutdown()
                    except:
                        pass

    def _is_connection_valid(self, producer) -> bool:
        """检查连接是否有效"""
        try:
            # 这里可以添加连接健康检查逻辑
            # 目前简单返回 True，实际使用时可以发送心跳消息
            return True
        except:
            return False

    def close(self):
        """关闭连接池"""
        self._closed = True

        # 关闭池中的所有连接
        while not self._pool.empty():
            try:
                producer = self._pool.get_nowait()
                producer.shutdown()
            except queue.Empty:
                break
            except Exception as e:
                logger.error(f"关闭连接失败: {e}")

        # 关闭活跃连接
        for producer in list(self._active_connections):
            try:
                producer.shutdown()
            except Exception as e:
                logger.error(f"关闭活跃连接失败: {e}")

        logger.info("连接池已关闭")

    def get_stats(self) -> Dict[str, Any]:
        """获取连接池统计信息"""
        return {
            "pool_size": self.pool_size,
            "available_connections": self._pool.qsize(),
            "active_connections": len(self._active_connections),
            "created_connections": self._created_count,
            "is_closed": self._closed
        }


class MessageProducer:
    """生产级消息发送者"""

    def __init__(self, config=None, connection_pool_size: int = 5):
        """
        初始化消息发送者

        Args:
            config: RocketMQ 配置
            connection_pool_size: 连接池大小
        """
        # 如果没有提供配置，使用默认配置
        if config is None:
            # config = get_config()  # 假设已实现
            config = self._get_default_config()

        self.config = config
        self.connection_pool = ConnectionPool(config, connection_pool_size)
        self.circuit_breaker = CircuitBreaker()

        # 性能监控
        self._stats = {
            "messages_sent": 0,
            "messages_failed": 0,
            "total_send_time": 0,
            "retry_count": 0
        }
        self._stats_lock = Lock()

        # 线程池用于异步发送
        self.executor = ThreadPoolExecutor(max_workers=20, thread_name_prefix="rocketmq-sender")

        # 健康检查
        self._health_check_event = Event()
        self._start_health_check()

        logger.info("MessageProducer 初始化完成")

    def _get_default_config(self):
        """获取默认配置（简化版，实际使用时应该使用配置管理器）"""
        from types import SimpleNamespace

        config = SimpleNamespace()
        config.name_server_address = "localhost:9876"
        config.producer_group = "DefaultProducerGroup"
        config.producer_max_message_size = 4 * 1024 * 1024
        config.producer_send_timeout = 3000
        config.producer_retry_times = 3
        config.access_key = None
        config.secret_key = None
        config.enable_health_check = True
        config.health_check_interval = 30

        return config

    def send_message(self, topic: str, body: Union[str, bytes, dict],
                     tags: Optional[str] = None, keys: Optional[str] = None,
                     properties: Optional[Dict[str, str]] = None,
                     priority: MessagePriority = MessagePriority.NORMAL,
                     delay_level: int = 0) -> SendResponse:
        """
        同步发送消息

        Args:
            topic: 主题名称
            body: 消息体
            tags: 消息标签
            keys: 消息键
            properties: 消息属性
            priority: 消息优先级
            delay_level: 延时级别（0-18）

        Returns:
            SendResponse: 发送响应
        """
        # 创建消息上下文
        message_ctx = MessageContext(
            message_id=str(uuid.uuid4()),
            topic=topic,
            tags=tags,
            keys=keys,
            body=body,
            properties=properties,
            priority=priority,
            delay_level=delay_level
        )

        return self._send_with_retry(message_ctx)

    def _send_with_retry(self, message_ctx: MessageContext) -> SendResponse:
        """带重试的发送消息"""
        start_time = time.time()
        last_exception = None

        for attempt in range(message_ctx.max_retries + 1):
            try:
                # 检查熔断器
                if not self.circuit_breaker.can_execute():
                    return SendResponse(
                        status=SendStatus.CIRCUIT_BREAKER_OPEN,
                        message_id=message_ctx.message_id,
                        error=Exception("熔断器开启"),
                        retry_count=attempt,
                        elapsed_time=time.time() - start_time
                    )

                # 执行发送
                response = self._do_send_message(message_ctx)

                # 记录成功
                self.circuit_breaker.record_success()
                self._update_stats(True, time.time() - start_time, attempt)

                return response

            except Exception as e:
                last_exception = e
                message_ctx.retry_count = attempt

                logger.warning(f"发送消息失败 (attempt {attempt + 1}/{message_ctx.max_retries + 1}): {e}")

                # 记录失败
                self.circuit_breaker.record_failure()

                # 如果不是最后一次重试，等待一段时间
                if attempt < message_ctx.max_retries:
                    delay = self._calculate_retry_delay(attempt)
                    time.sleep(delay)

        # 所有重试都失败了
        self._update_stats(False, time.time() - start_time, message_ctx.max_retries)

        return SendResponse(
            status=SendStatus.FAILED,
            message_id=message_ctx.message_id,
            error=last_exception,
            retry_count=message_ctx.max_retries,
            elapsed_time=time.time() - start_time
        )

    def _do_send_message(self, message_ctx: MessageContext) -> SendResponse:
        """实际发送消息的方法"""
        with self.connection_pool.get_connection() as producer:
            # 创建消息对象
            message = Message(message_ctx.topic)

            if isinstance(message_ctx.body, str):
                message.set_body(message_ctx.body.encode('utf-8'))
            else:
                message.set_body(message_ctx.body)

            if message_ctx.tags:
                message.set_tags(message_ctx.tags)

            if message_ctx.keys:
                message.set_keys(message_ctx.keys)

            if message_ctx.properties:
                for key, value in message_ctx.properties.items():
                    message.set_property(key, str(value))

            # 设置延时级别
            if message_ctx.delay_level > 0:
                message.set_delay_time_level(message_ctx.delay_level)

            # 发送消息
            send_result = producer.send_sync(message)

            return SendResponse(
                status=SendStatus.SUCCESS,
                message_id=send_result.msg_id,
                queue_id=send_result.queue_id,
                queue_offset=send_result.queue_offset,
                retry_count=message_ctx.retry_count
            )

    def send_message_async(self, topic: str, body: Union[str, bytes, dict],
                           tags: Optional[str] = None, keys: Optional[str] = None,
                           properties: Optional[Dict[str, str]] = None,
                           callback: Optional[Callable[[SendResponse], None]] = None) -> asyncio.Future:
        """
        异步发送消息

        Args:
            topic: 主题名称
            body: 消息体
            tags: 消息标签
            keys: 消息键
            properties: 消息属性
            callback: 回调函数

        Returns:
            asyncio.Future: 异步结果
        """

        def _async_send():
            response = self.send_message(topic, body, tags, keys, properties)
            if callback:
                callback(response)
            return response

        return asyncio.get_event_loop().run_in_executor(self.executor, _async_send)

    def send_batch_messages(self, messages: List[MessageContext]) -> List[SendResponse]:
        """
        批量发送消息

        Args:
            messages: 消息列表

        Returns:
            List[SendResponse]: 发送响应列表
        """
        results = []

        # 使用线程池并发发送
        future_to_message = {}

        with ThreadPoolExecutor(max_workers=min(10, len(messages))) as executor:
            for msg_ctx in messages:
                future = executor.submit(self._send_with_retry, msg_ctx)
                future_to_message[future] = msg_ctx

            for future in as_completed(future_to_message):
                try:
                    response = future.result()
                    results.append(response)
                except Exception as e:
                    msg_ctx = future_to_message[future]
                    results.append(SendResponse(
                        status=SendStatus.FAILED,
                        message_id=msg_ctx.message_id,
                        error=e
                    ))

        return results

    def _calculate_retry_delay(self, attempt: int) -> float:
        """计算重试延迟时间"""
        # 指数退避 + 随机抖动
        base_delay = 2 ** attempt
        max_delay = 60  # 最大延迟60秒
        delay = min(base_delay, max_delay)

        # 添加随机抖动（±20%）
        jitter = delay * 0.2 * (random.random() * 2 - 1)
        return max(0.1, delay + jitter)

    def _update_stats(self, success: bool, elapsed_time: float, retry_count: int):
        """更新统计信息"""
        with self._stats_lock:
            if success:
                self._stats["messages_sent"] += 1
            else:
                self._stats["messages_failed"] += 1

            self._stats["total_send_time"] += elapsed_time
            self._stats["retry_count"] += retry_count

    def _start_health_check(self):
        """启动健康检查"""
        if not self.config.enable_health_check:
            return

        def health_check_loop():
            while not self._health_check_event.wait(self.config.health_check_interval):
                try:
                    self._perform_health_check()
                except Exception as e:
                    logger.error(f"健康检查失败: {e}")

        health_thread = threading.Thread(target=health_check_loop, daemon=True)
        health_thread.start()
        logger.info("健康检查线程已启动")

    def _perform_health_check(self):
        """执行健康检查"""
        try:
            # 发送一个测试消息
            test_topic = "HEALTH_CHECK_TOPIC"
            response = self.send_message(
                topic=test_topic,
                body="health_check",
                tags="health"
            )

            if response.status == SendStatus.SUCCESS:
                logger.debug("健康检查成功")
            else:
                logger.warning(f"健康检查失败: {response.error}")

        except Exception as e:
            logger.error(f"健康检查异常: {e}")

    def get_stats(self) -> Dict[str, Any]:
        """获取统计信息"""
        with self._stats_lock:
            stats = self._stats.copy()

        # 添加连接池统计
        pool_stats = self.connection_pool.get_stats()
        stats.update(pool_stats)

        # 计算平均发送时间
        if stats["messages_sent"] > 0:
            stats["avg_send_time"] = stats["total_send_time"] / stats["messages_sent"]
        else:
            stats["avg_send_time"] = 0

        # 计算成功率
        total_attempts = stats["messages_sent"] + stats["messages_failed"]
        if total_attempts > 0:
            stats["success_rate"] = stats["messages_sent"] / total_attempts
        else:
            stats["success_rate"] = 0

        return stats

    def close(self):
        """关闭生产者"""
        self._health_check_event.set()  # 停止健康检查
        self.executor.shutdown(wait=True)  # 关闭线程池
        self.connection_pool.close()  # 关闭连接池
        logger.info("MessageProducer 已关闭")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
